// Copyright 2024 The Google Research Authors.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

#ifndef MOLECULE_TOOLS_SMU_H_
#define MOLECULE_TOOLS_SMU_H_

#include <memory>
#include <vector>

#include "third_party/lilly_mol/Foundational/cmdline/cmdline.h"
#include "third_party/lilly_mol/Foundational/iwmisc/report_progress.h"
#include "third_party/lilly_mol/Foundational/iwstring/iw_stl_hash_map.h"
#include "third_party/lilly_mol/Foundational/iwstring/iw_stl_hash_set.h"
#include "third_party/lilly_mol/Molecule_Lib/molecule.h"

namespace smu {

struct SMU_Params {
  resizable_array_p<Atom> to_add;
  int max_atoms;
  bool non_aromatic_unique_smiles;
  bool oneg_only_added_to_npos;
};

struct MoleculeData {
  // For each atom, the max number of connections.
  std::vector<int> max_conn;
  // For each atom, the max number of bonds.
  std::vector<int> max_bonds;
};

// Results can be stored in a large or small memory form.
class SmuResults {
 public:
  SmuResults();
  ~SmuResults();

  bool Add(std::unique_ptr<Molecule>& mol, IW_STL_Hash_Set& already_produced,
           bool non_aromatic_unique_smiles);

  void set_output_separator(char s) { _separator = s; }

  // Report statistics about molecules examined.
  // If `single_line`, the report is on a single record.
  void Report(std::ostream& output, bool single_line) const;

  // Write smiles file to `output`.
  // How molecules containing charged atoms is goverened by settings
  // in `params1.
  int Write(const SMU_Params& params, std::ostream& output) const;
  // Write smiles for our molecules to `output`.
  // Each molecule written is assigned `next_id` which is incremented.
  // How molecules containing charged atoms is handled is goverened by settings
  // in `params`.
  // Returns the number of molecules written.
  int Write(int& next_id, const SMU_Params& params, std::ostream& output) const;

  // Generate a smiles of all current molecules with that have
  // `natoms` atoms.
  // This is useful at the start of each expansion.
  std::vector<IWString> CurrentSmiles(int natoms) const;

  // The number of molecules currently stored.
  size_t size() const;

  int InitialiseProgressReporting(Command_Line& cl, char opt, int verbose) {
    return _report_progress.initialise(cl, opt, verbose);
  }

 private:
  // Mapping from smiles to identifier.
  // Used for detecting duplicates and storing results.
  IW_STL_Hash_Map<IWString, int> _smiles_to_id;

  // The number of molecules pass in via the Add method.
  int _molecules_examined;

  // The number rejected for being already present here.
  int _duplicates_discarded;

  // The number rejected for being already generated by
  // another SmuResults - the `already_produced` paramter to Add.

  int _already_produced_discarded;

  // The number discarded for invalid valences.
  int _invalid_valence_rejected;

  char _separator;

  Report_Progress _report_progress;

  // Private functions;

  template <typename ID>
  bool _maybe_write_molecule(Molecule& m, const IWString& smiles, const ID& id,
                             const SMU_Params& params,
                             std::ostream& output) const;
};

// Generate all molecules that can be created by adding to `m`.
// SMU expansion is controlled by `params`.
// `already_produced` contains smiles of molecules already
// generated - presumably from different starting molecules.
// It is updated with the molecules generated here.
// Results will be // in `smu_results`.
void Expand(Molecule& starting_molecule, const SMU_Params& params,
            IW_STL_Hash_Set& already_produced, SmuResults& smu_results);

// Return true if formal charges in `m` are ok.
// Behaviour is controlled by `params`.
bool OkCharges(const Molecule& m, const SMU_Params& params);
}  // namespace smu

#endif  // MOLECULE_TOOLS_SMU_H_
